import numpy as np
import argparse
import os
os.environ["OMP_NUM_THREADS"] = "16"
def gen_golden_data(M, N):
    os.makedirs('./input', exist_ok=True)
    os.makedirs('./output', exist_ok=True)
    
    A_dtype = np.float16
    B_dtype = np.float16
    Compute_dtype = np.float32
    C_dtype = np.float16

    x1_gm = np.random.randint(-5, 5, (M, N)).astype(A_dtype)
    x2_gm = np.random.randint(-5, 5, (M, N)).astype(B_dtype)
    x1_c = np.ascontiguousarray(x1_gm)
    x2_c = np.ascontiguousarray(x2_gm)
    golden = (x1_c + x2_c).astype(C_dtype)
    # print(x1_c, x2_c, golden) py一般是没错的
    x1_gm.tofile('./input/A.bin')
    x2_gm.tofile('./input/B.bin')
    golden.tofile('./output/golden.bin')


if __name__ == '__main__':
    M = 1000
    N = 1000

    parser = argparse.ArgumentParser()
    parser.add_argument('M', action='store', type=int)
    parser.add_argument('N', action='store', type=int)

    args = parser.parse_args()

    M = args.M
    N = args.N

    gen_golden_data(M, N)
